import os
import argparse

"""
encode raw yuv videos
"""

parser = argparse.ArgumentParser()
parser.add_argument("--seq_folder", type = str)
parser.add_argument("-sw", "--seq_width", type = str)
parser.add_argument("-sh", "--seq_height", type = str)

parser.add_argument("-c", "--codec", type = str, default = "libx265")
parser.add_argument("--crf", type = str, default = "20")
parser.add_argument("--gop", type = str, default = "12")
parser.add_argument("--preset", type = str, default = "veryfast")
parser.add_argument("--tune", type = str, default = "zerolatency")

parser.add_argument("--save_folder", type = str, default = "./encoded")
parser.add_argument("--format", type = str, default = "mkv")

args = parser.parse_args()

seq_folder = args.seq_folder
sw, sh = args.seq_width, args.seq_height

codec, crf, gop = args.codec, args.crf, args.gop
preset, tune = args.preset, args.tune

save_folder = args.save_folder
format = args.format


if not os.path.exists(save_folder):
    os.makedirs(save_folder)


for seq in os.listdir(seq_folder):
    if not seq.endswith(".yuv"):
        continue

    input = os.path.join(seq_folder, seq)
    options = [
        "ffmpeg",

        "-pix_fmt", "yuv420p",
        "-s:v", f"{sw}x{sh}",
        "-i", input,

        "-c:v", codec,
        "-preset", preset,
        "-tune", tune,
        "-x265-params", f"crf={crf}:keyint={gop}:verbose=1:csv-log-level=1:csv={save_folder}/report/{seq[:-4]}_report.csv",

        f"{save_folder}/{seq[:-4]}.{format}"
    ]

    cmd = ""
    for opt in options:
        cmd += (opt + " ")

    print(cmd)
    os.system(cmd)

